from fastapi import FastAPI, UploadFile, File, HTTPException, Query
from fastapi.responses import JSONResponse, StreamingResponse
import cv2
import numpy as np
from io import BytesIO
import torch
from 逻辑代码 import CLAHEPreprocessor,EyeDiseasePredictor,EyeFundusProcessor,RetinalVesselSegmenter

# 初始化 FastAPI 应用
app = FastAPI(
    title="眼底图像处理与分析服务",
    description="提供眼底图像预处理、疾病预测、视杯视盘渗出物显示以及血管分割功能",
    version="0.0.5"
)

# 初始化各类处理器和预测器
clahe_preprocessor = CLAHEPreprocessor(clip_limit=2.0, tile_grid_size=(8, 8))
eye_fundus_processor = EyeFundusProcessor()
vessel_segmenter = RetinalVesselSegmenter()

# 初始化疾病预测器
left_model_path = './model/best.pt'
right_model_path = './model/best.pt'

eye_disease_predictor = EyeDiseasePredictor(left_model_path, right_model_path)

# 定义返回值模型
from pydantic import BaseModel, Field
from typing import List


class PredictionResult(BaseModel):
    predicted_class_id: int = Field(..., example=3, description="预测的类别 ID")
    confidence: float = Field(..., example=0.9981, description="预测的置信度")
    class_name: str = Field(..., example="glaucoma", description="预测的疾病名称")
    inference_time: float = Field(..., example=0.375, description="推理时间（秒）")
    probabilities: List[float] = Field(..., example=[0.001, 0.0001, 0.998, 0.0005], description="所有类别的预测概率")

class EyePredictionResponse(BaseModel):
    left_eye: PredictionResult = Field(..., description="左眼的预测结果")
    right_eye: PredictionResult = Field(..., description="右眼的预测结果")

# 预处理图像 API
@app.post("/process_image/")
async def process_images(image: UploadFile = File(...)):
    """
    处理上传的单张图像，应用 CLAHE 预处理

    请求方法
    POST /process_image/

    请求参数
    参数名	类型	是否必填	描述	默认值
    image	file	是	上传的图像文件（应为一张的图像）	无

    响应格式
    200 OK：成功返回处理后的图像流，格式为 image/jpeg。
    500 Internal Server Error：图像处理过程中发生错误。

    响应内容
    成功时，响应体包含处理后的 JPEG 图像数据。
    """
    try:
        # 读取上传的图像
        image_data = await image.read()

        # 将图像数据转换为 OpenCV 格式
        img = cv2.imdecode(np.frombuffer(image_data, np.uint8), cv2.IMREAD_COLOR)

        if img is None:
            return {"error": "无法读取图像文件"}

        # 应用 CLAHE 预处理
        processed_img = clahe_preprocessor(img)

        # 将处理后的图像转换为字节流
        _, buffer = cv2.imencode(".jpg", processed_img)

        # 返回处理后的图像作为流式响应
        return StreamingResponse(BytesIO(buffer.tobytes()), media_type="image/jpeg")

    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"message": f"Error processing images: {str(e)}"}
        )

# 疾病预测 API
@app.post("/predict", response_model=EyePredictionResponse,     responses={
        200: {
            "description": "成功返回预测结果",
            "content": {
                "application/json": {
                    "example": {
    "Left Eye": {
        "Predicted Class ID": 3,
        "Confidence": 0.99811851978302,
        "Class Name": "glaucoma",
        "Inference Time": 0.3575100898742676,
        "Probabilities": [
            0.0010801143944263458,
            5.690049874829128e-06,
            5.455171869783726e-09,
            0.99811851978302,
            0.0007614473579451442,
            3.417734114918858e-05,
            1.1992923987236281e-07,
            1.9942584472687486e-08
        ]
    },
    "Right Eye": {
        "Predicted Class ID": 4,
        "Confidence": 0.9985502362251282,
        "Class Name": "hypertension",
        "Inference Time": 0.06664299964904785,
        "Probabilities": [
            0.000534689286723733,
            1.6410145690315403e-05,
            3.202781329036952e-07,
            0.0008948679897002876,
            0.9985502362251282,
            1.5368717640740215e-06,
            1.6064153669503867e-06,
            2.8503546900537913e-07
        ]
    }
}
                }
            }
        },
        500: {
            "description": "内部服务器错误",
            "content": {
                "application/json": {
                    "example": {"detail": "无法解码图像，请确保上传的是有效的图像文件。"}
                }
            }
        }
    })
async def predict(
    left_image: UploadFile = File(...),
    right_image: UploadFile = File(...)
):
    """
    对上传的左右眼图像进行疾病预测

    请求方法
    POST /predict

    请求参数
    参数名	类型	是否必填	描述	默认值
    left_image	file	是	上传的左眼图像文件	无
    right_image	file	是	上传的右眼图像文件	无

    响应格式
    200 OK：成功返回预测结果，格式为 application/json。
    500 Internal Server Error：预测过程中发生错误。

    响应内容
    成功时，返回一个 JSON 对象，包含左右眼的预测结果，每个结果包含：
    - predicted_class_id：预测的类别 ID
    - confidence：预测的置信度
    - class_name：预测的疾病名称
    - inference_time：推理时间（秒）
    - probabilities：所有类别的预测概率列表
    """
    try:
        # 读取上传的图像文件
        left_image_content = await left_image.read()
        right_image_content = await right_image.read()

        # 将图像内容转换为 NumPy 数组
        left_image_array = np.frombuffer(left_image_content, dtype=np.uint8)
        right_image_array = np.frombuffer(right_image_content, dtype=np.uint8)

        # 解码图像
        left_image = cv2.imdecode(left_image_array, cv2.IMREAD_COLOR)
        right_image = cv2.imdecode(right_image_array, cv2.IMREAD_COLOR)

        if left_image is None or right_image is None:
            raise ValueError("无法解码图像，请确保上传的是有效的图像文件。")

        # 调用预测器
        result = eye_disease_predictor.diagnose(left_image, right_image)

        # 处理返回结果，确保所有内容都是 JSON 可序列化的
        processed_result = {
            "left_eye": {
                "predicted_class_id": int(result["Left Eye"]["result"]["Predicted Class ID"]),
                "confidence": float(result["Left Eye"]["result"]["Confidence"].item()) if isinstance(result["Left Eye"]["result"]["Confidence"], torch.Tensor) else float(result["Left Eye"]["result"]["Confidence"]),
                "class_name": result["Left Eye"]["result"]["Class Name"],
                "inference_time": float(result["Left Eye"]["result"]["Inference Time"]),
                "probabilities": result["Left Eye"]["result"]["Probabilities"].tolist() if isinstance(result["Left Eye"]["result"]["Probabilities"], np.ndarray) else result["Left Eye"]["result"]["Probabilities"]
            },
            "right_eye": {
                "predicted_class_id": int(result["Right Eye"]["result"]["Predicted Class ID"]),
                "confidence": float(result["Right Eye"]["result"]["Confidence"].item()) if isinstance(result["Right Eye"]["result"]["Confidence"], torch.Tensor) else float(result["Right Eye"]["result"]["Confidence"]),
                "class_name": result["Right Eye"]["result"]["Class Name"],
                "inference_time": float(result["Right Eye"]["result"]["Inference Time"]),
                "probabilities": result["Right Eye"]["result"]["Probabilities"].tolist() if isinstance(result["Right Eye"]["result"]["Probabilities"], np.ndarray) else result["Right Eye"]["result"]["Probabilities"]
            }
        }

        return processed_result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# 视杯视盘渗出物显示 API
@app.post("/show_optic_disc_cup_exudate/")
async def process_eye_fundus(
    file: UploadFile = File(...),
    overlay: bool = Query(True, description="是否输出叠加掩码的图片"),
    bbox: bool = Query(True, description="是否输出绘制边框的图片"),
    alpha: float = Query(0.4, description="叠加掩码时的透明度")
):
    """
    处理眼底图像，检测并显示视盘、视杯以及渗出物区域

    请求方法
    POST /show_optic_disc_cup_exudate/

    请求参数
    参数名	类型	是否必填	描述	默认值
    file	file	是	上传的眼底图像文件	无
    overlay	boolean	否	是否输出叠加掩码的图片	True
    bbox	boolean	否	是否输出绘制边框的图片	True
    alpha	float	否	叠加掩码时的透明度，范围 0-1	0.4

    响应格式
    200 OK：成功返回处理后的图像流，格式为 multipart/x-mixed-replace; boundary=frame，每个图像部分以 --frame 分隔。
    400 Bad Request：图像解码失败。
    500 Internal Server Error：图像处理过程中发生错误。

    响应内容
    成功时，响应体包含多个 JPEG 图像数据，按顺序为：
    - 叠加视盘和渗出物掩码的图像（绿色透明叠加）
    - 叠加视杯掩码的图像（绿色透明叠加）
    - 绘制视盘和渗出物边框的图像（绿色边框）
    - 绘制视杯边框的图像（绿色边框）
    """
    # 读取上传的文件内容
    contents = await file.read()

    # 将文件内容转换为 OpenCV 图像格式
    nparr = np.frombuffer(contents, np.uint8)
    image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    if image is None:
        raise HTTPException(status_code=400, detail="Failed to decode image.")

    # 处理图像
    try:
        image_clahe, disc_mask, cup_mask, hist_green, hist_red = eye_fundus_processor.process_image_from_array(image)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")

    # 创建结果图像
    results = eye_fundus_processor.display_results(image_clahe, disc_mask, cup_mask, show_overlay=overlay, show_bbox=bbox, alpha=alpha)
    # print(f"生成的结果图片数量: {len(results)}")  # 调试信息

    # 将所有结果图像转换为 JPEG 格式并返回
    output_images = []
    for i, result in enumerate(results):
        _, buffer = cv2.imencode('.jpg', result)
        output_images.append(buffer.tobytes())

    # 返回所有结果图像
    async def generate():
        for img_data in output_images:
            yield img_data

    return StreamingResponse(generate(), media_type="multipart/x-mixed-replace; boundary=frame")

# 血管分割 API
@app.post("/segment_retinal_vessels")
async def segment_retinal_vessels(file: UploadFile = File(...)):
    """
    上传一张视网膜图像并返回血管分割结果

    请求方法
    POST /segment_retinal_vessels

    请求参数
    参数名	类型	是否必填	描述	默认值
    file	file	是	上传的视网膜图像文件	无

    响应格式
    200 OK：成功返回分割后的图像流，格式为 image/jpeg。
    500 Internal Server Error：图像处理过程中发生错误。

    响应内容
    成功时，响应体包含血管分割后的 JPEG 图像数据。
    """
    # 读取上传的文件
    contents = await file.read()
    nparr = np.frombuffer(contents, np.uint8)
    image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    if image is None:
        return {"error": "无法读取图像文件"}

    # 执行分割
    result_image = vessel_segmenter(image)

    result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB)

    # 将结果图像转换为字节流
    _, buffer = cv2.imencode(".jpg", result_image)
    return StreamingResponse(BytesIO(buffer.tobytes()), media_type="image/jpeg")

# 启动服务
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5000)
# 查看文档
# http://localhost:8000/docs#/
